package org.apache.lorry.zookeeper.cluster;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
@Slf4j
public class Worker {
    private String ipPort;

    private ZooKeeper zk;

    private boolean connected = false;

    private boolean expired = false;

    //本地缓存的cache
    private final ChildCache taskCache;

    private String name;

    private final Executor executor = Executors.newFixedThreadPool(3);

    public static final int SESSION_TIMEOUT = 10000;

    public Worker(String ipPort) {
        this.ipPort = ipPort;
        taskCache = new ChildCache();
    }

    public void startZK() throws Exception {
        CountDownLatch latch = new CountDownLatch(1);
        zk = new ZooKeeper(ipPort, SESSION_TIMEOUT, event -> {
            switch (event.getState()) {
                case SyncConnected:
                    connected = true;
                    latch.countDown();
                    break;
                case Disconnected:
                    connected = false;
                    break;
                case Expired:
                    connected = false;
                    expired = true;
                    break;
                default:
                    break;
            }
        });
        latch.await(SESSION_TIMEOUT, TimeUnit.MICROSECONDS);
        //创建worker父目录
        Stat stat = zk.exists(ZooConstant.WORKER_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.WORKER_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }

        stat = zk.exists(ZooConstant.ASSIGN_PATH, false);
        if (stat == null) {
            zk.create(ZooConstant.ASSIGN_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
    }

    public void register() {
        zk.create(ZooConstant.WORKER_PATH + "/", new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL, (rc, path, ctx, name) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    register();
                    break;
                case OK:
                    this.name = name.replaceAll(ZooConstant.WORKER_PATH + "/", "");
                    log.info("Registered successfully: {}", this.name);
                    bootstrap();
                    assignTask();
                    break;
                case NODEEXISTS:
                    log.info("Already registered: {}", this.name);
                    break;
                default:
                    System.out.println("Something went wrong: " + KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }, null);
    }

    public void bootstrap() {
        createAssignNode();
    }

    //创建这个任务的分配任务的节点
    private void createAssignNode() {
        zk.create(ZooConstant.ASSIGN_PATH + "/" + name, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT, (rc, path, ctx, name) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    createAssignNode();
                    break;
                case OK:
                    log.info("createAssignNode successfully: {}", path);
                    break;
                case NODEEXISTS:
                    log.warn("Already createAssignNode: {}", path);
                    break;
                default:
                    log.info("Something went wrong: {} , path: {}", KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }, null);
    }



    public void start() throws Exception {
        startZK();

        register();     //先注册，获取到对应的名称
    }

    //获取分配的对应任务
    private void assignTask() {
        zk.getChildren(ZooConstant.ASSIGN_PATH + "/" + name,event -> {
            if (event.getType().equals(Watcher.Event.EventType.NodeChildrenChanged)) {  //监听节点是否变化，变化进行处理
                assignTask();
            }
        },(rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    assignTask();
                    break;
                case OK:
                    if (children != null && !children.isEmpty()) {
                        List<String> newTaskList = taskCache.addChildDiff(children); //获取到新的task列表
                        if (CollectionUtils.isNotEmpty(newTaskList)) {
                            newTaskList.forEach(taskName -> {
                                executor.execute(new WorkTask(taskName, name, zk)); //提交任务
                            });
                        }
                    }
                    break;
                default:
                    log.error("getChildren failed: {}", KeeperException.create(KeeperException.Code.get(rc), path));
            }
        }, null);
    }

    public static void main(String[] args) throws Exception {
        Worker worker = new Worker("localhost:2181");
        worker.startZK();
        worker.register();
        TimeUnit.SECONDS.sleep(1000);
    }
}
